import asyncio
from datetime import datetime, timedelta

import cv2
import numpy as np
from aiortc import RTCPeerConnection, MediaStreamTrack
from aiortc.contrib.signaling import TcpSocketSignaling
from av import VideoFrame


class VideoReceiver:
    def __init__(self):
        self.track = None

    async def handle_track(self, track):
        print("处理帧")
        self.track = track
        frame_count = 0
        while True:
            try:
                print("等待帧...")
                frame = await asyncio.wait_for(track.recv(), timeout=5.0)
                frame_count += 1
                print(f"接收第 {frame_count} 帧")

                if isinstance(frame, VideoFrame):
                    print(f"帧类型: VideoFrame, pts: {frame.pts}, time_base: {frame.time_base}")
                    frame = frame.to_ndarray(format="rgb24")
                elif isinstance(frame, np.ndarray):
                    print(f"帧类型: numpy 数组")
                else:
                    print(f"意外的帧类型: {type(frame)}")
                    continue

                # 在帧上添加时间戳
                current_time = datetime.now()
                new_time = current_time - timedelta(seconds=55)
                timestamp = new_time.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
                cv2.putText(frame, timestamp, (10, frame.shape[0] - 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
                cv2.imwrite(f"imgs/received_frame_{frame_count}.jpg", frame)
                print(f"将帧 {frame_count} 保存至文件")
                # frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                cv2.imshow("Frame", frame)

                # 按 'q' 键退出程序
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            except asyncio.TimeoutError:
                print("等待帧超时，继续等待...")
            except Exception as e:
                print(f"handle_track 错误: {str(e)}")
                if "Connection" in str(e):
                    break
        print("退出 handle_track")


async def run(pc, signaling):
    await signaling.connect()

    print("等待发送方的 offer...")
    offer = await signaling.receive()
    print("接收 offer")
    await pc.setRemoteDescription(offer)
    print("设置远端描述")

    answer = await pc.createAnswer()
    print("创建 answer")
    await pc.setLocalDescription(answer)
    print("设置本地描述")

    await signaling.send(pc.localDescription)
    print("将 answer 发回给发送方")

    print("等待连接建立...")
    while pc.connectionState != "connected":
        await asyncio.sleep(0.1)

    print("连接建立，等待帧...")
    await asyncio.sleep(100)  # 等待接收帧 35 秒

    print("关闭连接")


async def main():
    signaling = TcpSocketSignaling("localhost", 9999)
    pc = RTCPeerConnection()

    video_receiver = VideoReceiver()

    @pc.on("track")
    def on_track(track):
        if isinstance(track, MediaStreamTrack):
            print(f"接收 {track.kind} 类型的轨道")
            asyncio.ensure_future(video_receiver.handle_track(track))

    @pc.on("datachannel")
    def on_datachannel(channel):
        print(f"数据通道 {channel.label} 已建立")

    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        print(f"连接状态是 {pc.connectionState}")
        if pc.connectionState == "connected":
            print("WebRTC 连接成功建立")

    try:
        await run(pc, signaling)
    except Exception as e:
        print(f"main 出现错误: {str(e)}")
    finally:
        print("关闭 peer 连接")
        await pc.close()


if __name__ == "__main__":
    asyncio.run(main())
